from transformers import BeitImageProcessor, BeitForImageClassification
from PIL import Image
import torch
image = Image.open('./lm.jfif')
processor = BeitImageProcessor.from_pretrained('./beit')
model = BeitForImageClassification.from_pretrained('./beit') #beit_base_patch16_224_pt22k.pth

inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])
